/*
 * AIEngine a new generation network intrusion detection system.
 *
 * Copyright (C) 2013-2023  Luis Campo Giralte
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Library General Public
 * License as published by the Free Software Foundation; either
 * version 2 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Library General Public License for more details.
 *
 * You should have received a copy of the GNU Library General Public
 * License along with this library; if not, write to the
 * Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
 * Boston, MA  02110-1301, USA.
 *
 * Written by Luis Campo Giralte <luis.camp0.2009@gmail.com>
 *
 */
#include "StackMobile.h"

namespace aiengine {

StackMobile::StackMobile() {

	name("Mobile Network Stack");

	// Add the Protocol objects
	addProtocol(eth, true);
	addProtocol(vlan, false);
	addProtocol(mpls, false);
	addProtocol(pppoe, false);
	addProtocol(ip_low_, true);
	addProtocol(udp_low_, true);
	addProtocol(gprs_, true);
	addProtocol(ip_high_, true);
	addProtocol(udp_high_, true);
	addProtocol(tcp_, true);
	addProtocol(icmp_, true);

	// Add the L7 protocols
	addProtocol(http);
	addProtocol(ssl);
	addProtocol(smtp);
	addProtocol(imap);
	addProtocol(pop);
	addProtocol(bitcoin);
	addProtocol(tcp_generic);
	addProtocol(freqs_tcp, false);
	addProtocol(dns);
	addProtocol(sip);
	addProtocol(ntp);
	addProtocol(snmp);
	addProtocol(ssdp);
	addProtocol(rtp);
	addProtocol(quic);
	addProtocol(udp_generic);
	addProtocol(freqs_udp, false);

        // Link the FlowCaches to their corresponding FlowManager for timeouts
        flow_table_udp_low_->setFlowCache(flow_cache_udp_low_);
        flow_table_udp_high_->setFlowCache(flow_cache_udp_high_);
        flow_table_tcp_->setFlowCache(flow_cache_tcp_);

	// Connect the Protocols with the Multiplexers
	ip_low_->setMultiplexer(mux_ip);
	mux_ip->setProtocol(static_cast<ProtocolPtr>(ip_low_));

	udp_low_->setMultiplexer(mux_udp_low_);
	mux_udp_low_->setProtocol(static_cast<ProtocolPtr>(udp_low_));
	ff_udp_low_->setProtocol(static_cast<ProtocolPtr>(udp_low_));

	gprs_->setFlowForwarder(ff_gprs_);
	gprs_->setMultiplexer(mux_gprs_);
	mux_gprs_->setProtocol(static_cast<ProtocolPtr>(gprs_));
	ff_gprs_->setProtocol(static_cast<ProtocolPtr>(gprs_));

        ip_high_->setMultiplexer(mux_ip_high_);
        mux_ip_high_->setProtocol(static_cast<ProtocolPtr>(ip_high_));

        udp_high_->setMultiplexer(mux_udp_high_);
        mux_udp_high_->setProtocol(static_cast<ProtocolPtr>(udp_high_));
        ff_udp_high_->setProtocol(static_cast<ProtocolPtr>(udp_high_));

	tcp_->setMultiplexer(mux_tcp_);
	mux_tcp_->setProtocol(static_cast<ProtocolPtr>(tcp_));
	ff_tcp_->setProtocol(static_cast<ProtocolPtr>(tcp_));

        icmp_->setMultiplexer(mux_icmp_);
        mux_icmp_->setProtocol(static_cast<ProtocolPtr>(icmp_));

	// Configure the multiplexers
	mux_eth->addUpMultiplexer(mux_ip);
	mux_ip->addDownMultiplexer(mux_eth);
	mux_ip->addUpMultiplexer(mux_udp_low_);
	mux_udp_low_->addDownMultiplexer(mux_ip);

	// configure the multiplexers of the second part
	mux_gprs_->addUpMultiplexer(mux_ip_high_);
        mux_ip_high_->addDownMultiplexer(mux_gprs_);
        mux_ip_high_->addUpMultiplexer(mux_icmp_);
	mux_icmp_->addDownMultiplexer(mux_ip_high_);
	mux_ip_high_->addUpMultiplexer(mux_tcp_);
	mux_tcp_->addDownMultiplexer(mux_ip_high_);
	mux_ip_high_->addUpMultiplexer(mux_udp_high_);
	mux_udp_high_->addDownMultiplexer(mux_ip_high_);

	// Connect the FlowManager and FlowCache
	tcp_->setFlowCache(flow_cache_tcp_);
	tcp_->setFlowManager(flow_table_tcp_);
	flow_table_tcp_->setProtocol(tcp_);

	udp_low_->setFlowCache(flow_cache_udp_low_);
	udp_low_->setFlowManager(flow_table_udp_low_);

	udp_high_->setFlowCache(flow_cache_udp_high_);
	udp_high_->setFlowManager(flow_table_udp_high_);
	flow_table_udp_high_->setProtocol(udp_high_);

        // Connect to upper layers the FlowManager
        http->setFlowManager(flow_table_tcp_);
        ssl->setFlowManager(flow_table_tcp_);
        smtp->setFlowManager(flow_table_tcp_);
        imap->setFlowManager(flow_table_tcp_);
        pop->setFlowManager(flow_table_tcp_);
        bitcoin->setFlowManager(flow_table_tcp_);
        dns->setFlowManager(flow_table_udp_high_);
        sip->setFlowManager(flow_table_udp_high_);
        ssdp->setFlowManager(flow_table_udp_high_);
        gprs_->setFlowManager(flow_table_udp_low_);

        freqs_tcp->setFlowManager(flow_table_tcp_);
        freqs_udp->setFlowManager(flow_table_udp_high_);

        // Connect the AnomalyManager with the protocols that may have anomalies
        ip_low_->setAnomalyManager(anomaly_);
        ip_high_->setAnomalyManager(anomaly_);
        tcp_->setAnomalyManager(anomaly_);
        udp_low_->setAnomalyManager(anomaly_);
        udp_high_->setAnomalyManager(anomaly_);

	// The low FlowManager have a 24 hours timeout to keep the Context on memory
        flow_table_udp_low_->setTimeout(86400);

	// Configure the FlowForwarders
	udp_low_->setFlowForwarder(ff_udp_low_);
	ff_udp_low_->addUpFlowForwarder(ff_gprs_);

	tcp_->setFlowForwarder(ff_tcp_);
	udp_high_->setFlowForwarder(ff_udp_high_);

	setTCPDefaultForwarder(ff_tcp_);
        setUDPDefaultForwarder(ff_udp_high_);

        std::ostringstream msg;
        msg << name() << " ready.";

        infoMessage(msg.str());

	setMode("full");
}

void StackMobile::showFlows(std::basic_ostream<char> &out, std::function<bool (const Flow&)> condition, int protocol) const {

        int total = flow_table_tcp_->getTotalFlows() + flow_table_udp_low_->getTotalFlows();
        total += flow_table_udp_high_->getTotalFlows();

        out << "Flows on memory " << total << std::endl;
        if (protocol == IPPROTO_TCP)
                flow_table_tcp_->showFlows(out, condition);
        else if (protocol == IPPROTO_UDP) {
                flow_table_udp_low_->showFlows(out, condition);
                flow_table_udp_high_->showFlows(out, condition);
	} else {
                flow_table_udp_low_->showFlows(out, condition);
                flow_table_tcp_->showFlows(out, condition);
                flow_table_udp_high_->showFlows(out, condition);
        }
}

void StackMobile::showFlows(Json &out, std::function<bool (const Flow&)> condition, int protocol) const {

        if (protocol == IPPROTO_TCP)
                flow_table_tcp_->showFlows(out, condition);
        else if (protocol == IPPROTO_UDP) {
                flow_table_udp_low_->showFlows(out, condition);
                flow_table_udp_high_->showFlows(out, condition);
	} else {
                flow_table_udp_low_->showFlows(out, condition);
                flow_table_tcp_->showFlows(out, condition);
                flow_table_udp_high_->showFlows(out, condition);
        }
}

void StackMobile::statistics(std::basic_ostream<char> &out) const {

        super_::statistics(out);
}

void StackMobile::setTotalTCPFlows(int value) {

        flow_cache_tcp_->create(value);
	tcp_->createTCPInfos(value);

	// The vast majority of the traffic of internet is HTTP
        // so create 75% of the value received for the http caches
	http->increaseAllocatedMemory(value * 0.75);

        // The 40% of the traffic is SSL
        ssl->increaseAllocatedMemory(value * 0.4);

        // 5% of the traffic could be SMTP/IMAP, im really positive :D
        smtp->increaseAllocatedMemory(value * 0.05);
        imap->increaseAllocatedMemory(value * 0.05);
        pop->increaseAllocatedMemory(value * 0.05);
        bitcoin->increaseAllocatedMemory(value * 0.05);
	freqs_tcp->increaseAllocatedMemory(16);
}

void StackMobile::setTotalUDPFlows(int value) {

	flow_cache_udp_high_->create(value);
        flow_cache_udp_low_->create(value/8);
        gprs_->increaseAllocatedMemory(value/8);
        dns->increaseAllocatedMemory(value / 2);
        sip->increaseAllocatedMemory(value * 0.2);
        ssdp->increaseAllocatedMemory(value * 0.2);
        quic->increaseAllocatedMemory(value * 0.2);
	freqs_udp->increaseAllocatedMemory(16);
}

int StackMobile::getTotalTCPFlows() const { return flow_cache_tcp_->getTotalFlows(); }

int StackMobile::getTotalUDPFlows() const { return flow_cache_udp_high_->getTotalFlows(); }

void StackMobile::setMode(const std::string &mode) {

        std::initializer_list<SharedPointer<FlowForwarder>> tcp_list {
                ff_tcp_, ff_http, ff_ssl, ff_smtp, ff_imap, ff_pop, ff_bitcoin,
                ff_tcp_generic, ff_tcp_freqs};
        std::initializer_list<SharedPointer<FlowForwarder>> udp_list {
                ff_udp_high_, ff_dns, ff_sip, ff_dhcp, ff_ntp, ff_snmp, ff_ssdp,
                ff_rtp, ff_quic, ff_dtls, ff_udp_generic, ff_udp_freqs};

        if (auto it = modes.find(mode); it != modes.end()) {
                std::ostringstream msg;

                disableFlowForwarders(tcp_list);
                disableFlowForwarders(udp_list);

                if ((*it).second == Mode::FULL) {

			operation_mode_ = "full";
                        msg << "Enable FullEngine mode on " << name();
                        enableFlowForwarders(tcp_list);
                        enableFlowForwarders(udp_list);

                } else if ((*it).second == Mode::FREQUENCIES) {

			operation_mode_ = "frequency";
                        msg << "Enable FrequencyEngine mode on " << name();
                        enableFlowForwarders({ff_tcp_, ff_tcp_freqs});
                        enableFlowForwarders({ff_udp_high_, ff_udp_freqs});

                } else if ((*it).second == Mode::NIDS) {

			operation_mode_ = "nids";
                        msg << "Enable NIDSEngine mode on " << name();
                        enableFlowForwarders({ff_tcp_, ff_tcp_generic});
                        enableFlowForwarders({ff_udp_high_, ff_udp_generic});
                }
                infoMessage(msg.str());
        }
}

void StackMobile::setFlowsTimeout(int timeout) {

        flow_table_tcp_->setTimeout(timeout);
        flow_table_udp_high_->setTimeout(timeout);
}

void StackMobile::setTCPRegexManager(const SharedPointer<RegexManager> &rm) {

        tcp_->setRegexManager(rm);
        tcp_generic->setRegexManager(rm);
	super_::setTCPRegexManager(rm);
}

void StackMobile::setUDPRegexManager(const SharedPointer<RegexManager> &rm) {

        udp_high_->setRegexManager(rm);
        udp_generic->setRegexManager(rm);
	super_::setUDPRegexManager(rm);
}

void StackMobile::setTCPIPSetManager(const SharedPointer<IPSetManager> &ipset_mng) {

        tcp_->setIPSetManager(ipset_mng);
        super_::setTCPIPSetManager(ipset_mng);
}

void StackMobile::setUDPIPSetManager(const SharedPointer<IPSetManager> &ipset_mng) {

        udp_high_->setIPSetManager(ipset_mng);
        super_::setUDPIPSetManager(ipset_mng);
}

#if defined(JAVA_BINDING)

void StackMobile::setTCPRegexManager(RegexManager *sig) {

        SharedPointer<RegexManager> rm;

        if (sig != nullptr)
                rm.reset(sig);

        setTCPRegexManager(rm);
}

void StackMobile::setUDPRegexManager(RegexManager *sig) {

        SharedPointer<RegexManager> rm;

        if (sig != nullptr)
                rm.reset(sig);

        setUDPRegexManager(rm);
}

void StackMobile::setTCPIPSetManager(IPSetManager *ipset_mng) {

        SharedPointer<IPSetManager> im;

        if (ipset_mng != nullptr)
                im.reset(ipset_mng);

        setTCPIPSetManager(im);
}

void StackMobile::setUDPIPSetManager(IPSetManager *ipset_mng) {

        SharedPointer<IPSetManager> im;

        if (ipset_mng != nullptr)
                im.reset(ipset_mng);

        setUDPIPSetManager(im);
}

#endif

std::tuple<Flow*, Flow*> StackMobile::getCurrentFlows() const {

        Flow *low_flow = udp_low_->getCurrentFlow();
        Flow *high_flow = nullptr;
	uint16_t proto = ip_high_->getProtocol();

        if (proto == IPPROTO_TCP)
                high_flow = tcp_->getCurrentFlow();
        else if (proto == IPPROTO_UDP)
                high_flow = udp_high_->getCurrentFlow();

#if GCC_VERSION < 50500
        return std::tuple<Flow*, Flow*>(low_flow, high_flow);
#else
        return {low_flow, high_flow};
#endif
}

SharedPointer<Flow> StackMobile::getFlow(const FlowSearchOptions &fsos) const {

	if (fsos.protocol == IPPROTO_UDP) {
		SharedPointer<Flow> flow = flow_table_udp_low_->find(fsos);
		if (!flow)
			flow = flow_table_udp_high_->find(fsos);

		return flow;
	} else
                return flow_table_tcp_->find(fsos);
}

} // namespace aiengine
